package burp.Utils;

import burp.IExtensionHelpers;
import burp.IParameter;
import burp.IRequestInfo;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

public class RequestUtils {

    /**
     * 更新 Content-Length 头
     */
    public static byte[] updateContentLength(
            IExtensionHelpers helpers,
            byte[] requestBytes,
            int newContentLength) {

        IRequestInfo requestInfo = helpers.analyzeRequest(requestBytes);
        List<String> headers = new ArrayList<>(requestInfo.getHeaders());

        // 替换或添加 Content-Length
        boolean found = false;
        for (int i = 0; i < headers.size(); i++) {
            if (headers.get(i).startsWith("Content-Length:")) {
                headers.set(i, "Content-Length: " + newContentLength);
                found = true;
                break;
            }
        }

        if (!found) {
            headers.add("Content-Length: " + newContentLength);
        }

        // 获取 body
        int bodyOffset = requestInfo.getBodyOffset();
        byte[] body = Arrays.copyOfRange(requestBytes, bodyOffset, requestBytes.length);

        // 构造请求
        return helpers.buildHttpMessage(headers, body);
    }

    /**
     * 从请求头中获取指定 header 的值
     */
    public static String getHeaderValue(List<String> headers, String headerName) {
        for (String header : headers) {
            if (header.startsWith(headerName + ":")) {
                return header.substring(headerName.length() + 1).trim();
            }
        }
        return null;
    }

    /**
     * Header 更新
     */
    public static byte[] updateHeader(
            IExtensionHelpers helpers,
            List<String> headers,
            byte[] request,
            String headerName,
            String newHeaderValue) {

        boolean headerExists = false;
        for (String header : headers) {
            if (header.startsWith(headerName + ":")) {
                headerExists = true;
                break;
            }
        }

        if (!headerExists && newHeaderValue != null) {
            return request;
        }

        // 更新 Header 列表
        for (int i = 0; i < headers.size(); i++) {
            if (headers.get(i).startsWith(headerName + ":")) {
                headers.set(i, headerName + ": " + newHeaderValue);
                break;
            }
        }

        // 解析原始请求以获取 bodyOffset
        IRequestInfo requestInfo = helpers.analyzeRequest(request);
        int bodyOffset = requestInfo.getBodyOffset();

        // 提取 body 内容
        byte[] body = Arrays.copyOfRange(request, bodyOffset, request.length);

        // 构造请求
        return helpers.buildHttpMessage(headers, body);
    }

    /**
     * 统一处理 URL 参数置空操作
     */
    public static byte[] nullifyUrlParameter(
            IExtensionHelpers helpers,
            IRequestInfo requestInfo,
            byte[] request,
            String paramName) {

        for (IParameter param : requestInfo.getParameters()) {
            if (param.getType() == IParameter.PARAM_URL && param.getName().equals(paramName)) {
                return helpers.updateParameter(request, helpers.buildParameter(param.getName(), "", param.getType()));
            }
        }
        return request;
    }

    /**
     * 统一处理 JSON Body 参数置空操作
     */
    public static byte[] nullifyJsonBodyParameter(
            IExtensionHelpers helpers,
            IRequestInfo requestInfo,
            byte[] request,
            String paramName) {

        int bodyOffset = requestInfo.getBodyOffset();
        byte[] jsonBodyBytes = Arrays.copyOfRange(request, bodyOffset, request.length);
        String jsonStr = helpers.bytesToString(jsonBodyBytes);

        Map<String, Object> jsonMap = JsonUtils.parseJsonBody(jsonStr);
        if (jsonMap == null || !jsonMap.containsKey(paramName)) {
            return request;
        }

        // 置空指定字段
        jsonMap.put(paramName, "");
        String updatedJson = JsonUtils.mapToJson(jsonMap);
        byte[] updatedBody = helpers.stringToBytes(updatedJson);

        // 构造请求
        return helpers.buildHttpMessage(requestInfo.getHeaders(), updatedBody);
    }

    /**
     * 置空 Form 表单类型的 Body 参数
     */
    public static byte[] nullifyFormBodyParameter(
            IExtensionHelpers helpers,
            IRequestInfo requestInfo,
            byte[] request,
            String paramName) {

        for (IParameter param : requestInfo.getParameters()) {
            if (param.getType() == IParameter.PARAM_BODY && param.getName().equals(paramName)) {
                try {
                    byte[] newRequest = helpers.updateParameter(
                            request,
                            helpers.buildParameter(param.getName(), "", IParameter.PARAM_BODY)
                    );
                    if (Arrays.equals(newRequest, request)) {
                        LoggerUtils.warn("Parameter not changed.");
                        return request;
                    }
                    return newRequest;
                } catch (Exception e) {
                    LoggerUtils.error("Failed to update parameter: ", e);
                }
            }
        }
        LoggerUtils.warn("Parameter not found.");
        return request;
    }

    /**
     * 置空 Cookie 参数
     */
    public static byte[] nullifyCookieParameter(
            IExtensionHelpers helpers,
            List<String> headers,
            byte[] request,
            String cookieName) {

        String cookieHeader = getHeaderValue(headers, "Cookie");
        if (cookieHeader == null) {
            return request;
        }

        List<String> updatedCookies = new ArrayList<>();
        for (String cookie : cookieHeader.split(";")) {
            String trimmed = cookie.trim();
            if (!trimmed.startsWith(cookieName + "=")) {
                updatedCookies.add(trimmed);
            } else {
                updatedCookies.add(cookieName + "="); // 清空值
            }
        }

        String newCookieHeader = String.join("; ", updatedCookies);

        // 更新或添加新的 Cookie 头
        boolean replaced = false;
        for (int i = 0; i < headers.size(); i++) {
            if (headers.get(i).startsWith("Cookie:")) {
                headers.set(i, "Cookie: " + newCookieHeader);
                replaced = true;
                break;
            }
        }
        if (!replaced) {
            headers.add("Cookie: " + newCookieHeader);
        }

        return updateHeader(helpers, headers, request, "Cookie", newCookieHeader);
    }

    /**
     * 统一处理参数置空操作
     */
    public static byte[] nullifyParameter(
            IExtensionHelpers helpers,
            IRequestInfo requestInfo,
            byte[] originalRequest,
            String paramName,
            boolean isUrlParam,
            boolean isHeader,
            boolean isCookie) {

        List<String> headers = new ArrayList<>(requestInfo.getHeaders());

        String contentType = getHeaderValue(headers, "Content-Type");
        Boolean isJson = contentType != null && contentType.contains("application/json");
        Boolean isFormPost = contentType != null && contentType.contains("application/x-www-form-urlencoded");

        byte[] modifiedRequest = originalRequest;
        if (isHeader || isCookie || isUrlParam || isJson || isFormPost) {
            modifiedRequest = Arrays.copyOf(originalRequest, originalRequest.length);
        }

        if (isHeader) {
            modifiedRequest = updateHeader(helpers, headers, originalRequest, paramName, "");
        } else if (isCookie) {
            modifiedRequest = nullifyCookieParameter(helpers, headers, originalRequest, paramName);
        } else if (isUrlParam) {
            modifiedRequest = nullifyUrlParameter(helpers, requestInfo, modifiedRequest, paramName);
        } else {
            if (isJson) {
                modifiedRequest = nullifyJsonBodyParameter(helpers, requestInfo, modifiedRequest, paramName);
            } else if (isFormPost) {
                modifiedRequest = nullifyFormBodyParameter(helpers, requestInfo, modifiedRequest, paramName);
            } else {
                return originalRequest;
            }
        }

        if (isJson || isFormPost) {
            IRequestInfo updatedInfo = helpers.analyzeRequest(modifiedRequest);
            int updatedBodyOffset = updatedInfo.getBodyOffset();
            byte[] newBody = Arrays.copyOfRange(modifiedRequest, updatedBodyOffset, modifiedRequest.length);

            // 构造请求
            modifiedRequest = helpers.buildHttpMessage(headers, newBody);

            // 更新 Content-Length
            modifiedRequest = updateContentLength(helpers, modifiedRequest, newBody.length);
        }

        return modifiedRequest;
    }
}
